import cv2
import numpy as np
import os
import matplotlib.pyplot as plt


def split_image_into_exercises(image_path, output_folder='output') -> None:
    # 读取图像
    image = cv2.imread(image_path)
    if image is None:
        print("Error: Could not load image.")
        return

    # 创建输出文件夹
    import os
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # 转换为灰度图
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 二值化图像（使用Otsu阈值法，并反转颜色）
    _, binary = cv2.threshold(
        gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)

    # 形态学操作：膨胀以连接文本行
    kernel = np.ones((5, 5), np.uint8)  # 调整核大小以连接行内的文本
    dilated = cv2.dilate(binary, kernel, iterations=2)

    # 查找轮廓
    contours, _ = cv2.findContours(
        dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 获取所有边界框
    boxes = []
    for contour in contours:
        x, y, w, h = cv2.boundingRect(contour)
        boxes.append((x, y, w, h))

    # 按y坐标排序边界框（从上到下）
    boxes.sort(key=lambda b: b[1])

    # 计算平均行高和分组行
    line_heights = [h for _, _, _, h in boxes]
    avg_line_height = np.mean(line_heights) if line_heights else 0

    # 分组边界框到行：合并同一行的框
    lines = []
    current_line = []
    current_y = -1
    for box in boxes:
        x, y, w, h = box
        if current_y == -1:
            current_y = y
        # 如果当前框的y坐标与当前行的y坐标相差小于平均行高的一半，则视为同一行
        if abs(y - current_y) < avg_line_height / 2:
            current_line.append(box)
        else:
            if current_line:
                # 合并当前行中的所有框：取最小x，最小y，最大x+w，最大y+h
                min_x = min(b[0] for b in current_line)
                min_y = min(b[1] for b in current_line)
                max_x = max(b[0] + b[2] for b in current_line)
                max_y = max(b[1] + b[3] for b in current_line)
                lines.append((min_x, min_y, max_x - min_x, max_y - min_y))
            current_line = [box]
            current_y = y
    if current_line:
        min_x = min(b[0] for b in current_line)
        min_y = min(b[1] for b in current_line)
        max_x = max(b[0] + b[2] for b in current_line)
        max_y = max(b[1] + b[3] for b in current_line)
        lines.append((min_x, min_y, max_x - min_x, max_y - min_y))

    # 按y坐标排序行
    lines.sort(key=lambda l: l[1])

    # 计算行之间的垂直距离，并分组题目
    exercise_groups = []
    current_group = []
    for i in range(len(lines)):
        if not current_group:
            current_group.append(lines[i])
        else:
            # 计算当前行与上一行的底部之间的距离
            prev_bottom = current_group[-1][1] + current_group[-1][3]
            current_top = lines[i][1]
            gap = current_top - prev_bottom
            # 如果间隙大于平均行高的1.5倍，则开始新组
            if gap > avg_line_height * 1.5:
                exercise_groups.append(current_group)
                current_group = [lines[i]]
            else:
                current_group.append(lines[i])
    if current_group:
        exercise_groups.append(current_group)

    # 为每个题目组创建边界框
    exercises = []
    for group in exercise_groups:
        min_x = min(b[0] for b in group)
        min_y = min(b[1] for b in group)
        max_x = max(b[0] + b[2] for b in group)
        max_y = max(b[1] + b[3] for b in group)
        exercises.append((min_x, min_y, max_x - min_x, max_y - min_y))

    # 裁剪和保存每个题目图片
    for i, (x, y, w, h) in enumerate(exercises):
        # 扩展边界一点，以确保完整捕获文本
        padding = 5
        x = max(0, x - padding)
        y = max(0, y - padding)
        w = min(image.shape[1] - x, w + 2 * padding)
        h = min(image.shape[0] - y, h + 2 * padding)
        crop_img = image[y:y+h, x:x+w]
        output_path = os.path.join(output_folder, f'exercise_{i+1}.png')
        cv2.imwrite(output_path, crop_img)
        print(f'Saved: {output_path}')


def process_and_display_steps(image_path, output_folder='output') -> None:
    """
    处理数学练习题图片并显示每一步结果
    """
    # 读取图像
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError("无法读取图像，请检查路径")

    # 创建输出目录
    os.makedirs(output_folder, exist_ok=True)

    # 步骤1: 显示原始图像
    plt.figure(figsize=(12, 8))
    plt.subplot(231)
    plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    plt.title('1. 原始图像')
    plt.axis('off')
    cv2.imwrite(os.path.join(output_folder, '01_original.jpg'), img)

    # 转换为灰度图
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    plt.subplot(232)
    plt.imshow(gray, cmap='gray')
    plt.title('2. 灰度图像')
    plt.axis('off')
    cv2.imwrite(os.path.join(output_folder, '02_gray.jpg'), gray)

    # 二值化处理
    _, binary = cv2.threshold(gray, 180, 255, cv2.THRESH_BINARY_INV)
    plt.subplot(233)
    plt.imshow(binary, cmap='gray')
    plt.title('3. 二值化图像（反转）')
    plt.axis('off')
    cv2.imwrite(os.path.join(output_folder, '03_binary.jpg'), binary)

    # 形态学操作 - 连接文本行
    kernel = np.ones((3, 30), np.uint8)
    dilated = cv2.dilate(binary, kernel, iterations=2)
    plt.subplot(234)
    plt.imshow(dilated, cmap='gray')
    plt.title('4. 膨胀操作后')
    plt.axis('off')
    cv2.imwrite(os.path.join(output_folder, '04_dilated.jpg'), dilated)

    # 查找轮廓
    contours, _ = cv2.findContours(
        dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 绘制轮廓
    contour_img = img.copy()
    cv2.drawContours(contour_img, contours, -1, (0, 255, 0), 2)
    plt.subplot(235)
    plt.imshow(cv2.cvtColor(contour_img, cv2.COLOR_BGR2RGB))
    plt.title('5. 检测到的轮廓')
    plt.axis('off')
    cv2.imwrite(os.path.join(output_folder, '05_contours.jpg'), contour_img)

    # 获取所有文本块的边界框
    text_blocks = []
    for contour in contours:
        x, y, w, h = cv2.boundingRect(contour)
        if w > 50 and h > 10:  # 过滤掉太小的区域
            text_blocks.append((x, y, w, h))

    # 按y坐标排序
    text_blocks.sort(key=lambda b: b[1])

    # 绘制边界框
    bbox_img = img.copy()
    for i, (x, y, w, h) in enumerate(text_blocks):
        cv2.rectangle(bbox_img, (x, y), (x + w, y + h), (255, 0, 0), 2)
        cv2.putText(bbox_img, str(i+1), (x, y-5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)

    plt.subplot(236)
    plt.imshow(cv2.cvtColor(bbox_img, cv2.COLOR_BGR2RGB))
    plt.title('6. 文本块边界框')
    plt.axis('off')
    cv2.imwrite(os.path.join(output_folder, '06_bboxes.jpg'), bbox_img)

    plt.tight_layout()
    plt.savefig(os.path.join(output_folder, 'processing_steps.jpg'),
                dpi=300, bbox_inches='tight')
    plt.show()

    # 分组：识别题目之间的间隔
    exercises = []
    current_exercise = []
    prev_bottom = 0

    for i, (x, y, w, h) in enumerate(text_blocks):
        if i == 0:
            current_exercise.append((x, y, w, h))
            prev_bottom = y + h
            continue

        # 计算与上一个块的垂直距离
        vertical_gap = y - prev_bottom

        # 如果间隙较大，开始新题目
        if vertical_gap > 50:
            if current_exercise:
                exercises.append(current_exercise)
            current_exercise = [(x, y, w, h)]
        else:
            current_exercise.append((x, y, w, h))

        prev_bottom = y + h

    if current_exercise:
        exercises.append(current_exercise)

    # 绘制题目分组结果
    group_img = img.copy()
    colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0),
              (255, 0, 255), (0, 255, 255), (128, 0, 128)]

    for idx, exercise_blocks in enumerate(exercises):
        color = colors[idx % len(colors)]

        # 计算整个题目的边界
        min_x = min(block[0] for block in exercise_blocks)
        min_y = min(block[1] for block in exercise_blocks)
        max_x = max(block[0] + block[2] for block in exercise_blocks)
        max_y = max(block[1] + block[3] for block in exercise_blocks)

        # 绘制大边界框
        cv2.rectangle(group_img, (min_x, min_y), (max_x, max_y), color, 3)
        cv2.putText(group_img, f'Exercise {idx+1}', (min_x, min_y-10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)

    plt.figure(figsize=(12, 8))
    plt.imshow(cv2.cvtColor(group_img, cv2.COLOR_BGR2RGB))
    plt.title('7. 题目分组结果')
    plt.axis('off')
    cv2.imwrite(os.path.join(output_folder, '07_grouping.jpg'), group_img)
    plt.show()

    # 为每个题目创建边界框并裁剪
    for idx, exercise_blocks in enumerate(exercises):
        # 计算整个题目的边界
        min_x = min(block[0] for block in exercise_blocks)
        min_y = min(block[1] for block in exercise_blocks)
        max_x = max(block[0] + block[2] for block in exercise_blocks)
        max_y = max(block[1] + block[3] for block in exercise_blocks)

        # 添加边距
        margin = 15
        min_x = max(0, min_x - margin)
        min_y = max(0, min_y - margin)
        max_x = min(img.shape[1], max_x + margin)
        max_y = min(img.shape[0], max_y + margin)

        # 裁剪并保存
        exercise_img = img[min_y:max_y, min_x:max_x]
        output_path = os.path.join(output_folder, f'exercise_{idx+1:02d}.png')
        cv2.imwrite(output_path, exercise_img)
        print(f'已保存: {output_path}')

        # 显示裁剪结果
        plt.figure(figsize=(8, 6))
        plt.imshow(cv2.cvtColor(exercise_img, cv2.COLOR_BGR2RGB))
        plt.title(f'Exercise {idx+1}')
        plt.axis('off')
        plt.show()


# 使用示例
if __name__ == "__main__":
    image_path = "1.png"  # 替换为您的图片路径
    process_and_display_steps(image_path)
